{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "020c5f4f0d944322b10d562dd51f9265": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_74ba1ddfa8194e7c82ef49622561a3fb",
              "IPY_MODEL_4a316fb7b9af4281a6bbf44ba47ad07c",
              "IPY_MODEL_f9ca4e73b4d04996a0e9c015ac95f22b"
            ],
            "layout": "IPY_MODEL_d03f501d29df44e38187fe3c7cca93c8"
          }
        },
        "74ba1ddfa8194e7c82ef49622561a3fb": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_c5ed3f385fce46839dc632a06af6e7ae",
            "placeholder": "​",
            "style": "IPY_MODEL_3f2cf23d18e7417aa2dfe51e155001c6",
            "value": "100%"
          }
        },
        "4a316fb7b9af4281a6bbf44ba47ad07c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_bcd8df79a65c4c70b733988edb5349a2",
            "max": 3,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_6a4af54835d3497aabbfbcc4940bacb4",
            "value": 3
          }
        },
        "f9ca4e73b4d04996a0e9c015ac95f22b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_90fb11617d154599afbd0ab8929c8e31",
            "placeholder": "​",
            "style": "IPY_MODEL_cdf329d0f12e4b199ca20cb43226f791",
            "value": " 3/3 [00:00&lt;00:00, 164.87it/s]"
          }
        },
        "d03f501d29df44e38187fe3c7cca93c8": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "c5ed3f385fce46839dc632a06af6e7ae": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "3f2cf23d18e7417aa2dfe51e155001c6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "bcd8df79a65c4c70b733988edb5349a2": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "6a4af54835d3497aabbfbcc4940bacb4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "90fb11617d154599afbd0ab8929c8e31": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "cdf329d0f12e4b199ca20cb43226f791": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "d90cc429a72b4dd6b50acaf6120f884b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_5a8436aa77f249d8a730fd3f7c404a4a",
              "IPY_MODEL_bde57097b0634f429abfed8a0649c506",
              "IPY_MODEL_d5430e6ce18d4861aacdfa9f9f76829c"
            ],
            "layout": "IPY_MODEL_92c5781644824e3d943f8ec0ae6fe939"
          }
        },
        "5a8436aa77f249d8a730fd3f7c404a4a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_862b5a87d2504ae6a25596ebd78096b9",
            "placeholder": "​",
            "style": "IPY_MODEL_d77ed94b7fb94174824b92eda94838a2",
            "value": "100%"
          }
        },
        "bde57097b0634f429abfed8a0649c506": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_7bf9463b8ca64f3bb2a3c7864b4a6743",
            "max": 3,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_6e94bd3d410e4fa48a4331dffa8ff123",
            "value": 3
          }
        },
        "d5430e6ce18d4861aacdfa9f9f76829c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_9021deedd19040b6a60aee40de0a6988",
            "placeholder": "​",
            "style": "IPY_MODEL_76634187d92d4d3c9a5d90fb18547455",
            "value": " 3/3 [00:00&lt;00:00, 135.25it/s]"
          }
        },
        "92c5781644824e3d943f8ec0ae6fe939": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "862b5a87d2504ae6a25596ebd78096b9": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d77ed94b7fb94174824b92eda94838a2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "7bf9463b8ca64f3bb2a3c7864b4a6743": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "6e94bd3d410e4fa48a4331dffa8ff123": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "9021deedd19040b6a60aee40de0a6988": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "76634187d92d4d3c9a5d90fb18547455": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "os.environ['PYTORCH_MPS_HIGH_WATERMARK_RATIO'] = '0.0'"
      ],
      "metadata": {
        "id": "guQ7M5IeGqMS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Handling Long Sequence by Chunking\n",
        "## Plan\n",
        "1. We will split the entire text into sentences.\n",
        "2. Calculate the embedding of each sentence.\n",
        "3. Calculate the whole text embedding by averaging the embeddings of each sentence.\n",
        "4. Pass the whole text embedding to the fully connected layer for text classification."
      ],
      "metadata": {
        "id": "_wCSEQbkDHyB"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Data Preparation"
      ],
      "metadata": {
        "id": "IV01qsVJEa_y"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 67,
          "referenced_widgets": [
            "020c5f4f0d944322b10d562dd51f9265",
            "74ba1ddfa8194e7c82ef49622561a3fb",
            "4a316fb7b9af4281a6bbf44ba47ad07c",
            "f9ca4e73b4d04996a0e9c015ac95f22b",
            "d03f501d29df44e38187fe3c7cca93c8",
            "c5ed3f385fce46839dc632a06af6e7ae",
            "3f2cf23d18e7417aa2dfe51e155001c6",
            "bcd8df79a65c4c70b733988edb5349a2",
            "6a4af54835d3497aabbfbcc4940bacb4",
            "90fb11617d154599afbd0ab8929c8e31",
            "cdf329d0f12e4b199ca20cb43226f791"
          ]
        },
        "id": "i7gX7KMRHtYn",
        "outputId": "c4212635-4d42-4e3e-8c7f-43bd41e8f925"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Found cached dataset imdb (/Users/premtimsina/.cache/huggingface/datasets/imdb/plain_text/1.0.0/d613c88cf8fa3bab83b4ded3713f1f74830d1100e171db75bbddb80b3345c9c0)\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/3 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "020c5f4f0d944322b10d562dd51f9265"
            }
          },
          "metadata": {}
        }
      ],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "from torch.utils.data import DataLoader, Dataset, random_split\n",
        "from datasets import load_dataset\n",
        "import nltk\n",
        "from transformers import BertTokenizer, BertModel\n",
        "\n",
        "# IMDB Dataset preparation\n",
        "imdb_data = load_dataset('imdb')\n",
        "\n",
        "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')\n",
        "# tokenizer.pad_token = tokenizer.eos_token\n",
        "\n",
        "class IMDBDataset(Dataset):\n",
        "    def __init__(self, data, tokenizer, max_sentence_length=48):\n",
        "        self.data = data\n",
        "        self.tokenizer = tokenizer\n",
        "        self.max_sentence_length = max_sentence_length\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.data)\n",
        "\n",
        "    \n",
        "    def __getitem__(self, idx):\n",
        "        text = self.data[idx]['text']\n",
        "        label = self.data[idx]['label']\n",
        "        # splitting the text into sentences\n",
        "        sentences = nltk.sent_tokenize(text)\n",
        "        # input_ids is the list of `input_ids` for each sentence\n",
        "        input_ids = [tokenizer.encode(sentence, max_length=self.max_sentence_length, truncation=True, padding='max_length') for sentence in sentences]\n",
        "        # attention_masks is the list of `attention_masks` for each sentence\n",
        "        attention_masks = [[1 if token_id != tokenizer.pad_token_id else 0 for token_id in sentence] for sentence in input_ids]\n",
        "        return {'input_ids': torch.tensor(input_ids, dtype=torch.long),\n",
        "                'attention_mask': torch.tensor(attention_masks, dtype=torch.long),\n",
        "                'label': torch.tensor(label, dtype=torch.long)}\n",
        "\n",
        "# We are just training with 10% of overall dataset. I am doing training in my laptop (Mac M2), which has just 30 core GPU and 32 GB RAM\n",
        "train_dataset = IMDBDataset(imdb_data['train'], tokenizer)\n",
        "test_dataset = IMDBDataset(imdb_data['test'], tokenizer)\n",
        "\n",
        "train_10_percent = int(len(train_dataset) * 0.1)\n",
        "test_10_percent = int(len(test_dataset) * 0.1)\n",
        "\n",
        "# Calculate the number of samples for the remaining 90%\n",
        "train_90_percent = len(train_dataset) - train_10_percent\n",
        "test_90_percent = len(test_dataset) - test_10_percent\n",
        "\n",
        "# Split the datasets into 10% and 90%\n",
        "train_data_10_percent, _ = random_split(train_dataset, [train_10_percent, train_90_percent])\n",
        "test_data_10_percent, _ = random_split(test_dataset, [test_10_percent, test_90_percent])\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### pad_collate\n",
        "**What we are doing here?**\n",
        "If you choose a batch size of 4, Pytorch requires that the number of sentences in each sample for that batch should be the same. However, this is not the case in real life. In a batch of size 4, sample 1 could have 10 sentences, and sample 2 could have 13 sentences.\n",
        "\n",
        "Therefore, in the pad_collate function, we perform padding on the sentences. We calculate the maximum number of sentences in that batch and pad accordingly."
      ],
      "metadata": {
        "id": "MKTwdM8dF8_o"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "def pad_collate(batch):\n",
        "    # calculate the maximum number of sentences in that batch\n",
        "    max_num_sentences = max([item['input_ids'].shape[0] for item in batch])\n",
        "\n",
        "    input_ids_batch = []\n",
        "    attention_masks_batch = []\n",
        "\n",
        "    # looping through each sample through batch. if batch_size =4; the following code loops for 4 times\n",
        "    for item in batch:\n",
        "        num_sentences = item['input_ids'].shape[0]\n",
        "        pad_length = max_num_sentences - num_sentences\n",
        "        input_ids = torch.cat([item['input_ids'], torch.zeros(pad_length, item['input_ids'].shape[1], dtype=torch.long)], dim=0)\n",
        "        attention_mask = torch.cat([item['attention_mask'], torch.zeros(pad_length, item['attention_mask'].shape[1], dtype=torch.long)], dim=0)\n",
        "\n",
        "        input_ids_batch.append(input_ids)\n",
        "        attention_masks_batch.append(attention_mask)\n",
        "\n",
        "    # it is just converting list into tensor by stacking along the dim=0\n",
        "    input_ids_tensor = torch.stack(input_ids_batch, dim=0)\n",
        "    attention_masks_tensor = torch.stack(attention_masks_batch, dim=0)\n",
        "    labels_tensor = torch.tensor([item['label'] for item in batch], dtype=torch.long)\n",
        "\n",
        "    return {'input_ids': input_ids_tensor, 'attention_mask': attention_masks_tensor, 'label': labels_tensor}\n",
        "\n",
        "train_loader = DataLoader(train_data_10_percent, batch_size=4, shuffle=True, collate_fn=pad_collate)\n",
        "test_loader = DataLoader(test_data_10_percent, batch_size=4, shuffle=True, collate_fn=pad_collate)"
      ],
      "metadata": {
        "id": "7KZwGv-oEjrv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "item=next(iter(train_loader))\n"
      ],
      "metadata": {
        "id": "mcdzMtouIdkk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### We are just looking at the shape of data.\n",
        "1. This signifies that batch_size=4, max_sentences inthat batch=21, max_seq_len=48"
      ],
      "metadata": {
        "id": "NHTSka2nIYIX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print('input_ids', item['input_ids'].shape)\n",
        "print('attention_mask', item['attention_mask'].shape)\n",
        "print('label', item['label'].shape)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JPoEcXGPH85n",
        "outputId": "24ac661f-53a0-428f-98e9-c84ff65f0401"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "input_ids torch.Size([4, 21, 48])\n",
            "attention_mask torch.Size([4, 21, 48])\n",
            "label torch.Size([4])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(item['input_ids'])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "056snCedKKas",
        "outputId": "614f2cf0-65b5-4278-ab4f-7fe3370575ed"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([[[ 101, 1045, 2031,  ...,    0,    0,    0],\n",
            "         [ 101, 5735, 2001,  ...,    0,    0,    0],\n",
            "         [ 101, 2200, 3422,  ...,    0,    0,    0],\n",
            "         [ 101, 2748, 2009,  ...,    0,    0,    0],\n",
            "         [   0,    0,    0,  ...,    0,    0,    0],\n",
            "         [   0,    0,    0,  ...,    0,    0,    0]],\n",
            "\n",
            "        [[ 101, 2044, 2108,  ..., 1000, 3458,  102],\n",
            "         [ 101, 2349, 2000,  ..., 2521, 2682,  102],\n",
            "         [ 101, 2023, 3576,  ...,    0,    0,    0],\n",
            "         [ 101, 2009, 2471,  ...,    0,    0,    0],\n",
            "         [ 101, 2174, 1010,  ..., 2477, 1007,  102],\n",
            "         [   0,    0,    0,  ...,    0,    0,    0]],\n",
            "\n",
            "        [[ 101, 2054, 2001,  ...,    0,    0,    0],\n",
            "         [ 101, 2076, 1996,  ..., 2787, 2008,  102],\n",
            "         [ 101, 2306, 1017,  ...,    0,    0,    0],\n",
            "         [ 101, 1045, 2179,  ...,    0,    0,    0],\n",
            "         [ 101, 1996, 2047,  ...,    0,    0,    0],\n",
            "         [ 101, 2006, 2028,  ..., 2028, 2008,  102]],\n",
            "\n",
            "        [[ 101, 2023, 2003,  ...,    0,    0,    0],\n",
            "         [ 101, 1045, 2031,  ...,    0,    0,    0],\n",
            "         [ 101, 2174, 1010,  ...,    0,    0,    0],\n",
            "         [ 101, 2122, 4301,  ..., 2065, 2009,  102],\n",
            "         [ 101, 2021, 2005,  ...,    0,    0,    0],\n",
            "         [   0,    0,    0,  ...,    0,    0,    0]]])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Model Architecture\n",
        "1. forward(self, input_ids_list, attention_mask_list): This method defines the forward pass of the neural network. It takes two arguments: input_ids_list and attention_mask_list, which are lists of input IDs and attention masks for each input sentence, respectively.\n",
        "\n",
        "2. The nested loops iterate through the batches and sentences within the batches. For each sentence, the BERT model is applied, and the 'pooler_output' is extracted, which is a fixed-size representation of the input text. The pooled outputs for each sentence are then stacked into a single tensor and averaged along the first dimension.\n",
        "\n",
        "3. The batch_pooled_outputs list is then converted into a tensor using torch.stack().\n",
        "\n",
        "\n",
        "4. The final fully connected (linear) layer is applied to produce the classification logits.\n",
        "\n",
        "5. The FineTunedBertClassifier class is designed to handle variable-length input sequences by processing them sentence by sentence and averaging their pooled outputs. This approach can handle different numbers of sentences per batch and allows the model to be fine-tuned on a specific classification task."
      ],
      "metadata": {
        "id": "3GXnwuWN3Axi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class FineTunedBertClassifier(nn.Module):\n",
        "    def __init__(self, n_classes):\n",
        "        super(FineTunedBertClassifier, self).__init__()\n",
        "        self.bert = BertModel.from_pretrained('bert-base-uncased')\n",
        "        self.dropout = nn.Dropout(p=0.3)\n",
        "        self.fc = nn.Linear(self.bert.config.hidden_size, n_classes)\n",
        "\n",
        "    def forward(self, input_ids_list, attention_mask_list):\n",
        "        batch_pooled_outputs = []\n",
        "\n",
        "        # Loop through each sample in a batch\n",
        "        for batch_input_ids, batch_attention_mask in zip(input_ids_list, attention_mask_list):\n",
        "            pooled_outputs = []\n",
        "            \n",
        "            # Loop through each sentence within the sample\n",
        "            # pooler_output is the sentence representation: output['pooler_output'] is the sentence embedding in case of bert\n",
        "            for input_ids, attention_mask in zip(batch_input_ids, batch_attention_mask):\n",
        "                output = self.bert(input_ids=input_ids.unsqueeze(0), attention_mask=attention_mask.unsqueeze(0))\n",
        "                pooler_output = output['pooler_output']\n",
        "                pooled_outputs.append(pooler_output)\n",
        "            \n",
        "            # Stack the pooled outputs into a single tensor\n",
        "            pooled_outputs_tensor = torch.stack(pooled_outputs, dim=1)\n",
        "            \n",
        "            # Compute the mean along the first dimension (dim=1)\n",
        "            concatenated_output = torch.mean(pooled_outputs_tensor, dim=1)\n",
        "            batch_pooled_outputs.append(concatenated_output)\n",
        "\n",
        "        # Stack the batch_pooled_outputs into a single tensor\n",
        "        batch_pooled_outputs_tensor = torch.stack(batch_pooled_outputs, dim=0)\n",
        "        \n",
        "        output = self.dropout(batch_pooled_outputs_tensor)\n",
        "        \n",
        "        output = torch.squeeze(output, dim=1)\n",
        "\n",
        "        return self.fc(output)"
      ],
      "metadata": {
        "id": "mLn9cyUMN7pr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The above code may looks confusing. Thus, lets understand through the example. It's important to understand that the following example is just to explain the architecture of model. Our IMDB data_set dimension is diffrent than following example.\n",
        "\n",
        "1. Suppose, input_ids_list and attention_mask_list looks like below. \n",
        "```\n",
        "input_ids_list = [\n",
        "    [  # sample 1\n",
        "        [101, 2023, 2003, 1037, 2742, 102],  # Sentence 1: \"This is a sample.\"\n",
        "        [101, 1045, 2066, 2026, 3835, 102]   # Sentence 2: \"I like my coffee.\"\n",
        "    ],\n",
        "    [  # sample 2\n",
        "        [101, 1045, 2572, 1037, 3076, 102],  # Sentence 1: \"I am a student.\"\n",
        "        [101, 2026, 3899, 2003, 2307, 102]   # Sentence 2: \"My dog is great.\"\n",
        "    ]\n",
        "]\n",
        "```\n",
        "```\n",
        "attention_mask_list = [\n",
        "    [  # sample 1\n",
        "        [1, 1, 1, 1, 1, 1],  # Attention mask for sentence 1\n",
        "        [1, 1, 1, 1, 1, 1]   # Attention mask for sentence 2\n",
        "    ],\n",
        "    [  # sample 2\n",
        "        [1, 1, 1, 1, 1, 1],  # Attention mask for sentence 1\n",
        "        [1, 1, 1, 1, 1, 1]   # Attention mask for sentence 2\n",
        "    ]\n",
        "]\n",
        "```\n",
        "2. pooled Output looks like\n",
        "\n",
        "\n",
        "```\n",
        "pooled_outputs = [\n",
        "    [tensor_1_1],  # pooler_output for Batch 1, Sentence 1\n",
        "    [tensor_1_2]   # pooler_output for Batch 1, Sentence 2\n",
        "]\n",
        "\n",
        "```\n",
        "Here, tensor_1_1 and tensor_1_2 are tensors of shape (1, 768) representing the pooled outputs (fixed-size sentence embeddings) for the two sentences in the first batch.\n",
        "\n",
        "When you execute pooled_outputs_tensor = torch.stack(pooled_outputs, dim=1), it stacks the pooled output tensors along dimension 1. The resulting pooled_outputs_tensor would have the shape (1, 2, 768):\n",
        "\n",
        "```\n",
        "[\n",
        "    [\n",
        "        [tensor_1_1],  # pooler_output for Batch 1, Sentence 1\n",
        "        [tensor_1_2]   # pooler_output for Batch 1, Sentence 2\n",
        "    ]\n",
        "]\n",
        "\n",
        "```\n",
        "3. \n",
        "In the example, the pooled_outputs_tensor has a shape of (1, 2, 768), where the first dimension corresponds to the batch, the second dimension corresponds to the sentences within the batch, and the third dimension corresponds to the hidden size of the BERT model's output.\n",
        "\n",
        "The line concatenated_output = torch.mean(pooled_outputs_tensor, dim=1) computes the mean of the pooled_outputs_tensor along dimension 1 (i.e., the sentences dimension). By doing this, we are effectively averaging the sentence embeddings within each batch. The result is a single tensor representing the entire batch, with a shape of (1, 768).\n",
        "\n",
        "In this specific example, the concatenated_output tensor would be:\n",
        "```\n",
        "[\n",
        "    [tensor_mean_1]  # Mean of pooled_outputs for Batch 1 (shape: (1, 768))\n",
        "]\n",
        "\n",
        "```\n",
        "\n",
        "4. batch_pooled_outputs\n",
        "\n",
        "\n",
        "```\n",
        "batch_pooled_outputs = [\n",
        "    [tensor_mean_1],  # Averaged sentence representation for sample 1\n",
        "    [tensor_mean_2]   # Averaged sentence representation for sample 2\n",
        "]\n",
        "\n",
        "```\n",
        "\n",
        "5. Finally Batched Pooled Output is send to the fully connected layer of classification"
      ],
      "metadata": {
        "id": "EgHNC1L65FKw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "len(train_loader)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nPZcaUQvPORT",
        "outputId": "769582f6-e6df-4056-d3df-95dbad3f011e"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "625"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "from torch.utils.data import DataLoader, Dataset, random_split\n",
        "from datasets import load_dataset\n",
        "from accelerate import Accelerator\n",
        "from tqdm import tqdm\n",
        "\n",
        "accelerator = Accelerator()\n",
        "device = accelerator.device\n",
        "\n",
        "model = FineTunedBertClassifier(n_classes=2)\n",
        "optimizer = optim.Adam(model.parameters(), lr=1e-5)\n",
        "\n",
        "model, optimizer = accelerator.prepare(model, optimizer)\n",
        "\n",
        "criterion = nn.CrossEntropyLoss()\n",
        "\n",
        "train_loader, test_loader = accelerator.prepare(train_loader, test_loader)\n",
        "# We are saving the intermediate model after every 100 batches.\n",
        "# This is a good practice if you have to train the model for many epochs.\n",
        "# There could be system failures or other issues.\n",
        "# Thus, you can restart your training from where you left off.\n",
        "# Adjust the save interval based on your needs.\n",
        "save_interval=100\n",
        "\n",
        "# Training loop\n",
        "num_epochs = 3\n",
        "\n",
        "for epoch in range(num_epochs):\n",
        "    model.train()\n",
        "    train_loader_progress = tqdm(enumerate(train_loader), desc=f\"Epoch {epoch + 1}/{num_epochs}, Training\", total=len(train_loader))\n",
        "    for batch_idx, batch in train_loader_progress:\n",
        "        input_ids = batch['input_ids'].to(device)\n",
        "        attention_mask = batch['attention_mask'].to(device)\n",
        "        labels = batch['label'].to(device)\n",
        "\n",
        "        optimizer.zero_grad()\n",
        "        logits = model(input_ids, attention_mask)\n",
        "        loss = criterion(logits, labels)\n",
        "        accelerator.backward(loss)\n",
        "        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
        "\n",
        "        optimizer.step()\n",
        "\n",
        "        if batch_idx % 10 == 0:\n",
        "            train_loader_progress.set_postfix(loss=loss.item())\n",
        "        if (batch_idx + 1) % save_interval == 0:\n",
        "            model_save_path = f\"/Users/premtimsina/Documents/bpbbook/chapter6/model/AlbertTextClassifier_epoch{epoch + 1}_batch{batch_idx + 1}.pt\"\n",
        "            torch.save(model.state_dict(), model_save_path)\n",
        "    model.eval()\n",
        "    total_correct = 0\n",
        "    total_samples = 0\n",
        "    test_loader_progress = tqdm(test_loader, desc=f\"Epoch {epoch + 1}/{num_epochs}, Testing\")\n",
        "\n",
        "    with torch.no_grad():\n",
        "        for batch in test_loader_progress:\n",
        "            input_ids = batch['input_ids'].to(device)\n",
        "            attention_mask = batch['attention_mask'].to(device)\n",
        "            labels = batch['label'].to(device)\n",
        "\n",
        "            logits = model(input_ids, attention_mask)\n",
        "            _, preds = torch.max(logits, dim=1)\n",
        "            total_correct += (preds == labels).sum().item()\n",
        "            total_samples += labels.size(0)\n",
        "\n",
        "    accuracy = total_correct / total_samples\n",
        "    print(f\"Epoch {epoch + 1}/{num_epochs}, Accuracy: {accuracy:.4f}\")"
      ],
      "metadata": {
        "id": "IgQeyJWPIFKj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "![Screenshot 2023-04-15 at 5.20.11 PM.png]()"
      ],
      "metadata": {
        "id": "jUFJh2ZaM4iA"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Analysis\n",
        "This is awesome. With just 10% of the dataset and only 1 epoch, we obtained an accuracy of 87%. I think the maximum accuracy for the IMDB dataset is around 96%. You could train the above model with the full dataset and a larger number of epochs to get even better result."
      ],
      "metadata": {
        "id": "YepGP5lLNGRH"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Hierirchal Attention\n",
        "## Plan\n",
        "1. Hierarchical attention: This model uses a two-level hierarchical attention mechanism:\n",
        "  * Local attention: Applied to individual sentences in a document to create sentence representations.\n",
        "  * Global attention: Applied to sentence representations to create a document representation.\n"
      ],
      "metadata": {
        "id": "p7UZSYibzI7z"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Data preparation\n",
        "The data preparation is exactly similar to data preparation we did above."
      ],
      "metadata": {
        "id": "nx2PP9sYjRvr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "from torch.utils.data import DataLoader, Dataset, random_split\n",
        "from datasets import load_dataset\n",
        "import nltk\n",
        "from transformers import BertTokenizer, BertModel\n",
        "\n",
        "# IMDB Dataset preparation\n",
        "imdb_data = load_dataset('imdb')\n",
        "\n",
        "from transformers import AlbertTokenizer\n",
        "\n",
        "tokenizer = AlbertTokenizer.from_pretrained(\"albert-base-v2\")\n",
        "\n",
        "class IMDBDataset(Dataset):\n",
        "    def __init__(self, data, tokenizer, max_sentence_length=48):\n",
        "        self.data = data\n",
        "        self.tokenizer = tokenizer\n",
        "        self.max_sentence_length = max_sentence_length\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.data)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        text = self.data[idx]['text']\n",
        "        label = self.data[idx]['label']\n",
        "        sentences = nltk.sent_tokenize(text)\n",
        "        input_ids = [tokenizer.encode(sentence, max_length=self.max_sentence_length, truncation=True, padding='max_length') for sentence in sentences]\n",
        "        attention_masks = [[1 if token_id != tokenizer.pad_token_id else 0 for token_id in sentence] for sentence in input_ids]\n",
        "        return {'input_ids': torch.tensor(input_ids, dtype=torch.long),\n",
        "                'attention_mask': torch.tensor(attention_masks, dtype=torch.long),\n",
        "                'label': torch.tensor(label, dtype=torch.long)}\n",
        "\n",
        "train_dataset = IMDBDataset(imdb_data['train'], tokenizer)\n",
        "test_dataset = IMDBDataset(imdb_data['test'], tokenizer)\n",
        "\n",
        "train_5_percent = int(len(train_dataset) * 0.05)\n",
        "test_5_percent = int(len(test_dataset) * 0.05)\n",
        "\n",
        "# Calculate the number of samples for the remaining 90%\n",
        "train_95_percent = len(train_dataset) - train_5_percent\n",
        "test_95_percent = len(test_dataset) - test_5_percent\n",
        "\n",
        "# Split the datasets into 10% and 90%\n",
        "train_data_5_percent, _ = random_split(train_dataset, [train_5_percent, train_95_percent])\n",
        "test_data_5_percent, _ = random_split(test_dataset, [test_5_percent, test_95_percent])\n",
        "\n",
        "def pad_collate(batch):\n",
        "    max_num_sentences = max([item['input_ids'].shape[0] for item in batch])\n",
        "\n",
        "    input_ids_batch = []\n",
        "    attention_masks_batch = []\n",
        "\n",
        "    for item in batch:\n",
        "        num_sentences = item['input_ids'].shape[0]\n",
        "        pad_length = max_num_sentences - num_sentences\n",
        "        input_ids = torch.cat([item['input_ids'], torch.zeros(pad_length, item['input_ids'].shape[1], dtype=torch.long)], dim=0)\n",
        "        attention_mask = torch.cat([item['attention_mask'], torch.zeros(pad_length, item['attention_mask'].shape[1], dtype=torch.long)], dim=0)\n",
        "\n",
        "        input_ids_batch.append(input_ids)\n",
        "        attention_masks_batch.append(attention_mask)\n",
        "\n",
        "    input_ids_tensor = torch.stack(input_ids_batch, dim=0)\n",
        "    attention_masks_tensor = torch.stack(attention_masks_batch, dim=0)\n",
        "    labels_tensor = torch.tensor([item['label'] for item in batch], dtype=torch.long)\n",
        "\n",
        "    return {'input_ids': input_ids_tensor, 'attention_mask': attention_masks_tensor, 'label': labels_tensor}\n",
        "\n",
        "train_loader = DataLoader(train_data_5_percent, batch_size=2, shuffle=True, collate_fn=pad_collate)\n",
        "test_loader = DataLoader(test_data_5_percent, batch_size=2, shuffle=True, collate_fn=pad_collate)\n",
        "\n"
      ],
      "metadata": {
        "id": "b7GUVas9_DuF",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 67,
          "referenced_widgets": [
            "d90cc429a72b4dd6b50acaf6120f884b",
            "5a8436aa77f249d8a730fd3f7c404a4a",
            "bde57097b0634f429abfed8a0649c506",
            "d5430e6ce18d4861aacdfa9f9f76829c",
            "92c5781644824e3d943f8ec0ae6fe939",
            "862b5a87d2504ae6a25596ebd78096b9",
            "d77ed94b7fb94174824b92eda94838a2",
            "7bf9463b8ca64f3bb2a3c7864b4a6743",
            "6e94bd3d410e4fa48a4331dffa8ff123",
            "9021deedd19040b6a60aee40de0a6988",
            "76634187d92d4d3c9a5d90fb18547455"
          ]
        },
        "outputId": "29fa9b5f-acc5-4fac-83e2-6dbcc23c3a09"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Found cached dataset imdb (/Users/premtimsina/.cache/huggingface/datasets/imdb/plain_text/1.0.0/d613c88cf8fa3bab83b4ded3713f1f74830d1100e171db75bbddb80b3345c9c0)\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/3 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "d90cc429a72b4dd6b50acaf6120f884b"
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "len(train_loader)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bf8pgtXH0jEW",
        "outputId": "4fedcc7a-8066-40e4-beae-6e5aaea3df70"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "625"
            ]
          },
          "metadata": {},
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 2.2. Model Architecture"
      ],
      "metadata": {
        "id": "eCFcPLlgjiWK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import AlbertModel\n",
        "\n",
        "class AlbertTextClassifier(nn.Module):\n",
        "    def __init__(self, num_classes):\n",
        "        super(AlbertTextClassifier, self).__init__()\n",
        "        self.albert = AlbertModel.from_pretrained(\"albert-base-v2\")\n",
        "        self.attention = nn.Linear(self.albert.config.hidden_size, 1)\n",
        "        self.classifier = nn.Linear(self.albert.config.hidden_size, num_classes)\n",
        "        self.dropout = nn.Dropout(0.1)\n",
        "\n",
        "    def forward(self, input_ids, attention_mask):\n",
        "        batch_size, num_sentences, seq_len = input_ids.shape\n",
        "        input_ids = input_ids.view(batch_size * num_sentences, seq_len)\n",
        "        attention_mask = attention_mask.view(batch_size * num_sentences, seq_len)\n",
        "\n",
        "        outputs = self.albert(input_ids, attention_mask=attention_mask)\n",
        "        hidden_states = outputs.last_hidden_state\n",
        "        attention_weights = torch.softmax(self.attention(hidden_states), dim=1)\n",
        "        sentence_representation = torch.sum(attention_weights * hidden_states, dim=1)\n",
        "        sentence_representation = sentence_representation.view(batch_size, num_sentences, -1)\n",
        "\n",
        "        doc_attention_weights = torch.softmax(self.attention(sentence_representation), dim=1)\n",
        "        doc_representation = torch.sum(doc_attention_weights * sentence_representation, dim=1)\n",
        "        doc_representation = self.dropout(doc_representation)  # Add dropout\n",
        "\n",
        "        logits = self.classifier(doc_representation)\n",
        "        return logits\n"
      ],
      "metadata": {
        "id": "Dlq0EVVxCvji"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "###The above code might be confusing: Let's explain through example:\n",
        "Assume we have the following input:\n",
        "\n",
        "  * 2 documents (batch_size = 2)\n",
        "  * Each document has 3 sentences (num_sentences = 3)\n",
        "  * Each sentence has 4 tokens (seq_len = 4)\n",
        "  * The input data would have the following dimensions:\n",
        "  ```\n",
        "  input_ids: (2, 3, 4)\n",
        "  attention_mask: (2, 3, 4)\n",
        "  ```\n",
        "Now, let's walk through the processing steps:\n",
        "\n",
        "1. Reshape input_ids and attention_mask:\n",
        "  ```\n",
        "  input_ids: (6, 4) [2 * 3 = 6, #since we have 2 documents with 3 sentences each]\n",
        "  attention_mask: (6, 4)\n",
        "  ```\n",
        "2. Pass input_ids and attention_mask to the ALBERT model:\n",
        "\n",
        "  * Obtain hidden_states with dimensions (6, 4, hidden_size)\n",
        "  * Compute attention weights for each token:\n",
        "\n",
        "3. Apply self.attention to hidden_states, resulting in a tensor of shape (6, 4, 1)\n",
        "4. Apply softmax to compute attention_weights: (6, 4, 1)\n",
        "5. Calculate sentence representations:\n",
        "\n",
        "  * Multiply attention_weights with hidden_states (element-wise): (6, 4, hidden_size)\n",
        "  * Sum along the sequence dimension (dim=1): (6, hidden_size)\n",
        "  * Reshape the tensor to (2, 3, hidden_size)\n",
        "\n",
        "6. Compute document-level attention weights:\n",
        "\n",
        "  * Apply self.attention to sentence_representation, resulting in a tensor of shape (2, 3, 1)\n",
        "  * Apply softmax to compute doc_attention_weights: (2, 3, 1)\n",
        "7. Calculate document representations:\n",
        "\n",
        "  * Multiply doc_attention_weights with sentence_representation (element-wise): (2, 3, hidden_size)\n",
        "  * Sum along the sentences dimension (dim=1): (2, hidden_size)\n",
        "  * Apply dropout to the document representations:\n",
        "\n",
        "8. doc_representation: (2, hidden_size)\n",
        "9. Pass the document representation through the classifier layer:\n"
      ],
      "metadata": {
        "id": "w4SuWMxNhO1s"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 2.3. Training loop"
      ],
      "metadata": {
        "id": "-opWvf-Ajnmn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "from torch.utils.data import DataLoader, Dataset, random_split\n",
        "from datasets import load_dataset\n",
        "from accelerate import Accelerator\n",
        "from tqdm import tqdm\n",
        "\n",
        "accelerator = Accelerator()\n",
        "device = accelerator.device\n",
        "\n",
        "model = AlbertTextClassifier(num_classes=2)\n",
        "optimizer = optim.Adam(model.parameters(), lr=1e-5)\n",
        "\n",
        "model, optimizer = accelerator.prepare(model, optimizer)\n",
        "\n",
        "criterion = nn.CrossEntropyLoss()\n",
        "\n",
        "train_loader, test_loader = accelerator.prepare(train_loader, test_loader)\n",
        "save_interval=100\n",
        "\n",
        "# Training loop\n",
        "num_epochs = 1\n",
        "\n",
        "for epoch in range(num_epochs):\n",
        "    model.train()\n",
        "    train_loader_progress = tqdm(enumerate(train_loader), desc=f\"Epoch {epoch + 1}/{num_epochs}, Training\", total=len(train_loader))\n",
        "    for batch_idx, batch in train_loader_progress:\n",
        "        input_ids = batch['input_ids'].to(device)\n",
        "        attention_mask = batch['attention_mask'].to(device)\n",
        "        labels = batch['label'].to(device)\n",
        "\n",
        "        optimizer.zero_grad()\n",
        "        logits = model(input_ids, attention_mask)\n",
        "        loss = criterion(logits, labels)\n",
        "        accelerator.backward(loss)\n",
        "        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
        "\n",
        "        optimizer.step()\n",
        "\n",
        "        if batch_idx % 10 == 0:\n",
        "            train_loader_progress.set_postfix(loss=loss.item())\n",
        "        if (batch_idx + 1) % save_interval == 0:\n",
        "            model_save_path = f\"/Users/premtimsina/Documents/bpbbook/chapter6/model/AlbertTextClassifier_epoch{epoch + 1}_batch{batch_idx + 1}.pt\"\n",
        "            torch.save(model.state_dict(), model_save_path)\n",
        "    model.eval()\n",
        "    total_correct = 0\n",
        "    total_samples = 0\n",
        "    test_loader_progress = tqdm(test_loader, desc=f\"Epoch {epoch + 1}/{num_epochs}, Testing\")\n",
        "\n",
        "    with torch.no_grad():\n",
        "        for batch in test_loader_progress:\n",
        "            input_ids = batch['input_ids'].to(device)\n",
        "            attention_mask = batch['attention_mask'].to(device)\n",
        "            labels = batch['label'].to(device)\n",
        "\n",
        "            logits = model(input_ids, attention_mask)\n",
        "            _, preds = torch.max(logits, dim=1)\n",
        "            total_correct += (preds == labels).sum().item()\n",
        "            total_samples += labels.size(0)\n",
        "\n",
        "    accuracy = total_correct / total_samples\n",
        "    print(f\"Epoch {epoch + 1}/{num_epochs}, Accuracy: {accuracy:.4f}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Pau8b1Cmzmwv",
        "outputId": "372dcae5-a685-428f-b297-0c9a32426543"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Some weights of the model checkpoint at albert-base-v2 were not used when initializing AlbertModel: ['predictions.dense.bias', 'predictions.dense.weight', 'predictions.bias', 'predictions.LayerNorm.bias', 'predictions.decoder.weight', 'predictions.decoder.bias', 'predictions.LayerNorm.weight']\n",
            "- This IS expected if you are initializing AlbertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
            "- This IS NOT expected if you are initializing AlbertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
            "Epoch 1/1, Training: 100%|█████| 625/625 [02:27<00:00,  4.24it/s, loss=0.000264]\n",
            "Epoch 1/1, Testing: 100%|█████████████████████| 625/625 [00:44<00:00, 14.15it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 1/1, Accuracy: 0.8688\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "It's awesome, we obtained 86% accuracy which is similar to document chuncking approach"
      ],
      "metadata": {
        "id": "h21R5y2vjsze"
      }
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "-tvIHsxb1Cu9"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}